package com.zlm.gateway.config;

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.*;
import org.springframework.http.HttpMethod;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.CharBuffer;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @Author: 郑
 *
 */
@Slf4j
public class CacheBodyGlobalFilter implements Ordered, GlobalFilter {
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        long startTime = System.currentTimeMillis();
        ServerHttpRequest request = exchange.getRequest();
        //如果是post请求,取出请求体
        String requestBodyStr = HttpMethod.POST.equals(request.getMethod()) ? resolveBodyFromRequest(request) : null;
        //获取请求信息
        String url = request.getPath().toString();
        Map<String, String> paramMap = request.getQueryParams().toSingleValueMap();
        log.info("请求地址:【{}】\n请求参数:GET【{}】\nPOST:【{}】\n", url, JSON.toJSONString(paramMap), requestBodyStr);


        //获取响应数据
        ServerHttpResponseDecorator decorator = printResponseBody(exchange,startTime);

        return chain.filter(exchange.mutate().request(request).response(decorator).build());
    }

    /**
     * 获取请求体
     * @param request
     * @return
     */
    private String resolveBodyFromRequest(ServerHttpRequest request) {
        Flux<DataBuffer> body = request.getBody();

        AtomicReference<String> bodyRef = new AtomicReference<>();
        body.subscribe(buffer -> {
            CharBuffer charBuffer = StandardCharsets.UTF_8.decode(buffer.asByteBuffer());
            DataBufferUtils.release(buffer);
            bodyRef.set(charBuffer.toString());
        });
        return bodyRef.get();
    }

    //打印响应数据
    private ServerHttpResponseDecorator printResponseBody(ServerWebExchange exchange,Long startTime) {
        ServerHttpResponse response = exchange.getResponse();
        DataBufferFactory bufferFactory = response.bufferFactory();
        ServerHttpResponseDecorator decorator = new ServerHttpResponseDecorator(response){
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                if (body instanceof Flux) {
                    Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
                    return super.writeWith(fluxBody.buffer().map(dataBuffer -> {
                        DefaultDataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
                        DefaultDataBuffer join = dataBufferFactory.join(dataBuffer);
                        byte[] content = new byte[join.readableByteCount()];
                        join.read(content);
                        //释放内存
                        DataBufferUtils.release(join);
                        String str = new String(content, StandardCharsets.UTF_8);

                        long time = System.currentTimeMillis() - startTime;
                        log.info("响应时间:【{}】\n响应数据:【{}】\n",time,str);
                        return bufferFactory.wrap(content);
                    }));

                }
                return super.writeWith(body);
            }
        };
        return decorator;
    }

    @Override
    public int getOrder() {

        return Ordered.HIGHEST_PRECEDENCE;
    }
}
